import glob
import random
import os
import sys
import numpy as np
from PIL import Image
import torch
from torch.utils.data import Dataset
import torch.nn.functional as F
import cv2
from torchvision.transforms import transforms
from utils.augmentations import horisontal_flip
from datasets.dataSetUtils import label_pad_norm, checklabel

import xml.etree.ElementTree as ET

def pad_to_square(img, pad_value):
    '''
    :param img:
    :param pad_value:
    :return: pad 左右上下pad
    '''
    c, h, w = img.shape
    dim_diff = np.abs(h - w)

    # (upper / left) padding and (lower / right) padding
    pad1, pad2 = int(dim_diff // 2), int(dim_diff - dim_diff // 2)
    pad = (0, 0, pad1, pad2) if h <= w else (pad1, pad2, 0, 0)
    img = F.pad(img, pad, "constant", value=pad_value)
    return img, pad

def resize(image, size):
    image = F.interpolate(image.unsqueeze(0), size=size, mode="nearest").squeeze(0)
    return image

def read_label(imagepath):
    '''
    :param label_path: 文件路径
    :return: box[x, y, w, h]
    '''
    label_path = os.path.join(os.path.dirname(imagepath) + "_label", os.path.basename(imagepath)[:-3] + "txt")
    targets = None
    if os.path.exists(label_path):
        boxes = torch.from_numpy(np.loadtxt(label_path).reshape(-1, 5))
    else:
        return targets
    return boxes

def get_element(root, childElementName):
    '''
    input：
        @root: 根节点
        @childElementName: 字节点tag名称
    output：
        @elements:根节点下第一个符合的子元素对象
    '''
    element = root.find(childElementName)
    return element

def check_path(data_path, dir):
    if not os.path.exists(os.path.join(data_path, dir)):
        os.makedirs(os.path.join(data_path, dir))


class listDataset(Dataset):
    def __init__(self, list_path, img_size=416, augment=True, multiscale=True, normalized_labels=True):
        '''
        :param list_path: 数据集的路径，这个文件是一个 txt 文件，里面每一行是一个图片的路径
        '''
        with open(list_path, 'r') as file:
            self.img_files = file.readlines()
            self.img_files = [x.strip() for x in self.img_files]

        if len(self.img_files) > 20:
            self.img_files = self.img_files[:20]

        self.label_files = []
        for img in self.img_files:
            label = os.path.dirname(img) + "_label"
            self.label_files.append(os.path.join(label, os.path.basename(img)[:-3] + "txt"))

        self.img_size = img_size
        self.max_objects = len(self.img_files)
        self.augment = augment
        self.multiscale = multiscale
        self.normalized_labels = normalized_labels
        self.min_size = self.img_size - 3 * 32
        self.max_size = self.img_size + 3 * 32
        self.batch_count = 0
        self.cls_map = {}

    def __getitem__(self, index):

        img_path = self.img_files[index].rstrip()

        # Extract image as PyTorch tensor
        img = transforms.ToTensor()(Image.open(img_path).convert('RGB'))
        # Handle images with less than three channels
        if len(img.shape) != 3:
            img = img.unsqueeze(0)
            img = img.expand((3, img.shape[1:]))

        _, h, w = img.shape
        h_factor, w_factor = (h, w) if self.normalized_labels else (1, 1)

        # pad = (0, 0, 0, 0)
        img, pad = pad_to_square(img, 0)
        _, padded_h, padded_w = img.shape

        #coco的label

        targets_coco = read_label(img_path)

        targets = label_pad_norm(targets_coco, pad, (padded_w, padded_h))

        checklabel(targets[:, 1:], img)

        # Apply augmentations
        if self.augment:
            if np.random.random() < 0.5:
                img, targets = horisontal_flip(img, targets)

        return img_path, img, targets

    def collate_fn(self, batch):
        paths, imgs, targets = list(zip(*batch))

        targets = [boxes for boxes in targets if boxes is not None]

        for i, boxes in enumerate(targets):
            boxes[:, 0] = i

        targets = torch.cat(targets, 0)

        # Selects new image size every tenth batch
        if self.multiscale and self.batch_count % 10 == 0:
            self.img_size = random.choice(range(self.min_size, self.max_size + 1, 32))

        # Resize images to input shape
        imgs = torch.stack([resize(img, self.img_size) for img in imgs])
        self.batch_count += 1
        return paths, imgs, targets

    def __len__(self):
        return len(self.img_files)

    def make_class_map(self, classnames):
        self.cls_map.clear()
        for id, name in enumerate(classnames):
            self.cls_map[name] = id
